"""
Case Type   : DiskANN向量检索
Case Name   : diskann索引逻辑备份恢复功能正常
Create At   : 2025-11-18
Owner       : sungang14
Description :
    1、创建table插入数据
    2、创建diskann索引
    3、创建用户并授权
    4、使用gs_dump进行逻辑备份
    5、向表中插入数据并查询
    6、使用gs_restore进行恢复
    7、查询数据
    8、清理环境
Expect      :
    1、开启成功
    2、创建成功，数据插入成功
    3、创建索引失败,报错diskann不支持极致RTO
    4、备份成功
    5、还原参数成功
    6、恢复成功
    7、索引正常，查询结果为备份之前的记录数
    8、删除用户、表和文件成功
History     :
"""

import os
import time
import unittest

from testcase.utils.CommonSH import CommonSH
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger
from yat.test import Node
from yat.test import macro

logger = Logger()
Primary_SH = CommonSH('PrimaryDbUser')


@unittest.skipIf(3 != Primary_SH.get_node_num(), '非1+2环境不执行')
class SlowSQLTestCase(unittest.TestCase):
    def setUp(self):
        logger.info(f'-----{os.path.basename(__file__)} start-----')
        self.constant = Constant()
        self.comsh = CommonSH('PrimaryDbUser')
        self.user_node = Node('PrimaryDbUser')
        status = self.comsh.get_db_cluster_status('status')
        self.assertTrue(status)
        self.table_name = f't_func_diskann_index0047'
        self.index_name = f'idx_func_diskann_index0047'
        self.prepare_data_00 = f"drop table if exists {self.table_name};"
        self.prepare_data_01 = f'''
                                create table {self.table_name} (id bigserial primary key, val1 vector(5));
                                declare
                                 vector_dim int := 5;
                                 vector_value float[];
                                 record_num int := 1000;
                                 i int;
                                 j int;
                                begin
                                 for j in 1..record_num loop
                                 for i in 1..vector_dim loop
                                 vector_value[i] := random()*1000+1;
                                 end loop;
                                 insert into {self.table_name} (val1) values (vector_value);
                                 end loop;
                                end;
                                '''
        self.prepare_data_02 = f"drop index if exists {self.index_name};" \
                               f"alter table {self.table_name} set (parallel_workers=32);" \
                               f"create index {self.index_name} on {self.table_name} using diskann(val1 vector_l2_ops);"
        self.dump_path = os.path.join(macro.DB_INSTANCE_PATH, 'dumpcase0047.tar')
        self.prepare_data_04 = f"source {macro.DB_ENV_PATH}; rm -rf {self.dump_path}; gs_dump -f {self.dump_path} -p {self.user_node.db_port} --include-depend-objs -t {self.table_name} {self.user_node.db_name}  -F t"
        self.prepare_data_05 = f"insert into {self.table_name} values(1001, '[2,2,2,2,2]');"
        self.prepare_data_06 = f"set enable_seqscan=off;" \
                               f"set enable_indexscan=on;" \
                               f"select * from {self.table_name} order by val1 <-> '[2,2,2,2,2]' limit 10;"
        self.prepare_data_07 = f"source {macro.DB_ENV_PATH};gs_restore {self.dump_path} -p {self.user_node.db_port} -d {self.user_node.db_name}"

    def test_slowsql(self):
        text = '--step1:创建table插入数据;expect:创建成功，数据插入成功--'
        logger.info(text)
        data_cmd = self.comsh.execut_db_sql(self.prepare_data_00)
        logger.info(data_cmd)
        self.assertIn(self.constant.TABLE_DROP_SUCCESS, data_cmd, '执行失败' + text)
        data_cmd = self.comsh.execut_db_sql(self.prepare_data_01)
        logger.info(data_cmd)
        self.assertIn(self.constant.TABLE_CREATE_SUCCESS, data_cmd, '执行失败' + text)
        self.assertIn(f'ANONYMOUS BLOCK EXECUTE', data_cmd, '执行失败' + text)

        text = '--step2:创建diskann索引;expect:创建失败，报错diskann不支持极致RTO--'
        logger.info(text)
        data_cmd = self.comsh.execut_db_sql(self.prepare_data_02)
        logger.info(data_cmd)
        self.assertIn(self.constant.DROP_INDEX_SUCCESS_MSG, data_cmd, '执行失败' + text)
        self.assertIn(self.constant.ALTER_TABLE_MSG, data_cmd, '执行失败' + text)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS, data_cmd, '执行失败' + text)

        text = '--step4：使用gs_dump进行逻辑备份;expect:备份成功--'
        logger.info(text)
        logger.info(self.prepare_data_04)
        msg5 = self.user_node.sh(self.prepare_data_04).result()
        logger.info(msg5)
        self.assertIn('dump database {} successfully'.format(
            self.user_node.db_name), msg5.split('\n')[-2].strip())

        text = '--step5:向表中插入数据并查询; expect: 成功--'
        logger.info(text)
        data_cmd = self.comsh.execut_db_sql(self.prepare_data_05)
        logger.info(data_cmd)
        self.assertIn(f"INSERT 0 1", data_cmd, '执行失败' + text)
        data_cmd = self.comsh.execut_db_sql(self.prepare_data_06)
        logger.info(data_cmd)
        self.assertIn(f"1001 | [2,2,2,2,2]", data_cmd, '执行失败' + text)
        self.assertIn(f"10 rows", data_cmd, '执行失败' + text)

        text = '--step6: 使用gs_restore进行恢复；expect：成功--'
        logger.info(text)
        data_cmd = self.comsh.execut_db_sql(self.prepare_data_00)
        logger.info(data_cmd)
        self.assertIn(self.constant.TABLE_DROP_SUCCESS, data_cmd, '执行失败' + text)
        logger.info(self.prepare_data_07)
        msg4 = self.user_node.sh(self.prepare_data_07).result()
        logger.info(msg4)
        self.assertIn(self.constant.RESTORE_SUCCESS_MSG,
                      msg4.split('\n')[-2].strip())

        text = '--step7:查询数据；expect：索引正常，查询结果为备份之前的记录数；--'
        logger.info(text)
        data_cmd = self.comsh.execut_db_sql(self.prepare_data_06)
        logger.info(data_cmd)
        self.assertNotIn(f"1001 | [2,2,2,2,2]", data_cmd, '执行失败' + text)

    def tearDown(self):
        text = '--step8:清理环境，恢复默认值;expect:成功--'
        logger.info(text)
        data_cmd = self.comsh.execut_db_sql(self.prepare_data_00)
        logger.info(data_cmd)
        self.assertIn(self.constant.TABLE_DROP_SUCCESS, data_cmd, '执行失败' + text)
        self.user_node.sh(f"rm -rf {self.dump_path}")
        logger.info(f'-----{os.path.basename(__file__)} end-----')
